from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from a06_export import YOLOv8Pipeline
import cv2
import numpy as np
import uvicorn
from pathlib import Path
from typing import List, Dict
import json

"""
调用api
curl -X POST -F "file=@F:/TwWork/ODVS/data/1.jpeg" http://localhost:8000/predict
"""
class DetectionAPI:
    def __init__(self, config_path: str = "configs/default.yaml"):
        self.pipeline = YOLOv8Pipeline(config_path)
        self.pipeline.load_model()

        # 加载导出模型
        self.export_path = Path(self.pipeline.config['paths']['runs']['export']) / 'model.onnx'
        if not self.export_path.exists():
            self.export_path = self.pipeline.export_model('onnx')

        # 创建API应用
        self.app = FastAPI(
            title="YOLOv8 Detection API",
            description="API for object detection using YOLOv8",
            version="1.0.0"
        )

        # 添加路由
        self.app.post("/predict")(self.predict)
        self.app.get("/model-info")(self.get_model_info)

    async def predict(self, file: UploadFile = File(...)) -> Dict:
        """处理预测请求"""
        try:
            # 读取图像
            img_bytes = await file.read()
            img = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)
            if img is None:
                raise HTTPException(status_code=400, detail="Invalid image format")

            # 使用原始模型预测（或使用导出的ONNX模型）
            results = self.pipeline.model(img)[0]

            # 格式化结果
            return {
                "boxes": results.boxes.xyxy.tolist(),
                "scores": results.boxes.conf.tolist(),
                "class_ids": results.boxes.cls.tolist(),
                "class_names": [self.pipeline.class_names[int(i)] for i in results.boxes.cls]
            }

        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))

    async def get_model_info(self) -> Dict:
        """获取模型信息"""
        try:
            with open(self.pipeline.eval_dir / 'metrics.json') as f:
                metrics = json.load(f)

            return {
                "model_name": self.pipeline.config['training']['base']['name'],
                "classes": self.pipeline.class_names,
                "metrics": metrics
            }
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))


def run_api_service(host: str = "0.0.0.0", port: int = 8000):
    """启动API服务"""
    api = DetectionAPI()
    print(f"🚀 API server starting on http://{host}:{port}")
    uvicorn.run(api.app, host=host, port=port)


if __name__ == "__main__":
    run_api_service()